package org.example.test2;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

public class DistributedLock {
    // zookeeper 服务器地址
    private static final String CONNECT_STRING = "47.108.48.236:2182";
    // 会话超时时间（毫秒）
    private static final int SESSION_TIMEOUT = 300000;
    private ZooKeeper zooKeeper;
    // 加锁节点
    private String lockPath;
    // 等待连接建立
    private CountDownLatch connectLatch = new CountDownLatch(1);
    // 创建的临时顺序节点
    private String curNode;
    // 顺序节点前缀
    private static final String CREATE_NODE_PREFIX = "seq-";
    // 监听节点
    private String waitNode;
    private CountDownLatch waitLatch = new CountDownLatch(1);

    public DistributedLock(String lockPath) {
        this.lockPath = lockPath;
    }

    /**
     * 建立 zookeeper 连接
     * @throws InterruptedException
     * @throws IOException
     */
    public void connect() throws InterruptedException, IOException {
        zooKeeper = new ZooKeeper(CONNECT_STRING, SESSION_TIMEOUT, new Watcher() {
            @Override
            public void process(WatchedEvent watchedEvent) {
                if (watchedEvent.getState() == Event.KeeperState.SyncConnected) {
                    // 连接成功，释放阻塞
                    connectLatch.countDown();
                }
                // 若前一节点删除, 获取到锁
                if (watchedEvent.getType() == Event.EventType.NodeDeleted && watchedEvent.getPath().equals(waitNode)) {
                    System.out.println("节点" + waitNode +"删除, 获取到锁...");
                    waitLatch.countDown();
                }
            }
        });
        // 等待连接建立
        connectLatch.await();
        System.out.println("zookeeper 连接成功");
    }

    /**
     * 加锁
     * @throws InterruptedException
     * @throws KeeperException
     */
    public void zkLock() throws InterruptedException, KeeperException {
        // 目标节点是否存在
        if (zooKeeper.exists(lockPath, null) == null) {
            throw new RuntimeException("目标节点: " + lockPath + " 不存在!");
        }
        /**
         * 创建临时顺序节点
         * String path: 创建节点路径
         * byte[] data: 节点数据
         * List<ACL> acl: 访问控制列表（Access Control List）,定义谁可以访问该节点
         * CreateMode createMode: 创建节点类型
         */
        curNode = zooKeeper.create(lockPath + "/" + CREATE_NODE_PREFIX, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        // 获取 目标节点 下的节点
        List<String> children = zooKeeper.getChildren(lockPath, null);
        // 计算当前节点位置
        Collections.sort(children);
        String thisNode = curNode.substring((lockPath + "/").length());
        int curIndex = children.indexOf(thisNode);
        if (curIndex == 0) {
            // 获取到锁
            System.out.println(curNode + "获取到锁...");
        } else if (curIndex > 0) {
            // 监听前一节点
            waitNode = lockPath + "/" + children.get(curIndex - 1);
            Stat stat = zooKeeper.exists(waitNode, true);
            // 前一节点删除, 获取锁
            if (stat == null) {
                System.out.println("节点" + waitNode +"删除, 获取到锁...");
            } else {
                // 等待前一节点释放锁
                waitLatch.await();
            }
        } else {
            throw new RuntimeException("节点获取失败! children: " + children + " node: " + curNode);
        }
    }

    public void zkUnlock() throws InterruptedException, KeeperException {
        zooKeeper.delete(curNode, -1);
        System.out.println(curNode + " 删除, 释放锁...");
    }
}
